"""
At the command line, only need to run once to install the package via pip:

$ pip install --upgrade google-generativeai
"""

import os
import time
import json
import google.generativeai as genai
from markdown_process import extract_json_code_blocks
import socket

socket.setdefaulttimeout(180)


GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
if GOOGLE_API_KEY is None:
    raise ValueError("GOOGLE_API_KEY environment variable must be set.")
genai.configure(api_key=GOOGLE_API_KEY, transport="rest")
video_file_path = os.path.join(os.path.dirname(__file__), "../video/test.mp4")

print("Uploading file...")
video_file = genai.upload_file(path=video_file_path, display_name="test video")
print(f"Completed upload: {video_file.uri}")
while video_file.state.name == "PROCESSING":
    print("Waiting for video to be processed.")
    time.sleep(10)
    video_file = genai.get_file(video_file.name)

if video_file.state.name == "FAILED":
    raise ValueError(video_file.state.name)
print("Video processing complete: " + video_file.uri)

# # Set up the model
# generation_config = {
#     "temperature": 1,
#     "top_p": 0.95,
#     "top_k": 0,
#     "max_output_tokens": 8192,
# }
# safety_settings = [
#     {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
#     {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
#     {
#         "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
#         "threshold": "BLOCK_MEDIUM_AND_ABOVE",
#     },
#     {
#         "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
#         "threshold": "BLOCK_MEDIUM_AND_ABOVE",
#     },
# ]
# # Set the model to Gemini 1.5 Pro.
# model = genai.GenerativeModel(
#     model_name="models/gemini-1.5-pro-latest",
#     generation_config=generation_config,
#     safety_settings=safety_settings,
# )
instruction = """Suppose you are a robot equiped with a camera navigating in the world. You can see the video recorded by the camera and output the timestamp when you arrive at and leave each scene. Extract all the meaningful objects you think you can use as a waypoint for navigation in each scene. 
Output as following json format:
```json
[
    {
        "scene": scene_name,
        "arrive_time": arrive_time,
        "leave_time": leave_time,
        "waypoints": {
            object_name: closest_time,
            ...
        }
    },
    ...
]
```
all time must be 'minutes:seconds' format or N/A meaning throughout the video.
"""
model = genai.GenerativeModel(
    model_name="models/gemini-1.5-pro-latest", system_instruction=instruction
)

# history = [
#     {
#         "role": "user",
#         "parts": [
#             '',
#         ],
#     },
#     {
#         "role": "model",
#         "parts": [
#             '',
#         ],
#     },
# ]
chat_session = model.start_chat(history=[])

response = chat_session.send_message(video_file)
# response = model.generate_content(
#     [prompt, video_file], request_options={"timeout": 600}
# )

# 打开文件并将 Python 对象写入文件
json_text = extract_json_code_blocks(response.text)
if len(json_text) > 0:
    json_dict = json.loads(json_text[0])
print("\ntime map as follow:")
print(json_dict)
print()
with open(os.path.join(os.path.dirname(__file__), "../json/time.json"), "w") as f:
    json.dump(obj=json_dict, fp=f)

genai.delete_file(video_file.name)
print(f"Deleted file {video_file.uri}")
